package com.example.base;

import com.example.entity.Customer;
import com.example.entity.CustomerJdbcStatement;
import com.example.entity.MysqlXaDataSourceSupplier;
import com.mysql.cj.jdbc.MysqlXADataSource;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.connector.base.DeliveryGuarantee;
import org.apache.flink.connector.jdbc.JdbcConnectionOptions;
import org.apache.flink.connector.jdbc.JdbcExactlyOnceOptions;
import org.apache.flink.connector.jdbc.JdbcExecutionOptions;
import org.apache.flink.connector.jdbc.core.datastream.sink.JdbcSink;
import org.apache.flink.connector.jdbc.core.datastream.source.JdbcSource;
import org.apache.flink.connector.jdbc.core.datastream.source.reader.extractor.ResultExtractor;
import org.apache.flink.connector.jdbc.datasource.connections.SimpleJdbcConnectionProvider;
import org.apache.flink.connector.jdbc.datasource.connections.xa.SimpleXaConnectionProvider;
import org.apache.flink.connector.jdbc.datasource.statements.JdbcQueryStatement;
import org.apache.flink.connector.kafka.source.KafkaSource;
import org.apache.flink.connector.kafka.source.enumerator.initializer.OffsetsInitializer;
import org.apache.flink.formats.json.JsonDeserializationSchema;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;

import java.sql.PreparedStatement;
import java.sql.SQLException;
import java.util.Optional;

/**
 * Author wangJinLong
 * Date 2025/8/21 08:26
 **/
public class FlinkStreamJdbc {
    public static void main(String[] args) throws Exception {
        exactlyOnce();
    }

    private static void exactlyOnce() throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);
        env.enableCheckpointing(10000);

        KafkaSource<Customer> source = KafkaSource.<Customer>builder()
                .setBootstrapServers("106.54.174.109:19092")
                .setTopics("data-user")
                .setGroupId("my-group1")
                .setStartingOffsets(OffsetsInitializer.latest())
                .setValueOnlyDeserializer(new JsonDeserializationSchema<>(Customer.class))
                .build();

        DataStreamSource<Customer> dataStreamSource =
                env.fromSource(source, WatermarkStrategy.noWatermarks(), "Kafka Source");

        dataStreamSource.print();

        // 创建 XA 数据源
        MysqlXADataSource xaDataSource = new MysqlXADataSource();
        xaDataSource.setUrl("jdbc:mysql://106.54.174.109:13306/my_cdc");
        xaDataSource.setUser("root");
        xaDataSource.setPassword("mysql@123456");

        MysqlXaDataSourceSupplier supplier = new MysqlXaDataSourceSupplier(xaDataSource);

        dataStreamSource.sinkTo(
                new JdbcSink<>(
                        DeliveryGuarantee.EXACTLY_ONCE,
                        SimpleXaConnectionProvider.from(supplier, null),
                        JdbcExecutionOptions.builder()
                                .withMaxRetries(0) // 必须写0
                                .withBatchSize(10)
                                .withBatchIntervalMs(5000)
                                .build(),
                        JdbcExactlyOnceOptions.builder()
                                .withTimeoutSec(Optional.of(60))
                                .build(),
                        new CustomerJdbcStatement()
                )
        );

        env.execute();
    }

    private static void atLeastOnceSink() throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);

        KafkaSource<Customer> source = KafkaSource.<Customer>builder()
                .setBootstrapServers("106.54.174.109:19092")
                .setTopics("data-user")
                .setGroupId("my-group1")
                .setStartingOffsets(OffsetsInitializer.latest())
                .setValueOnlyDeserializer(new JsonDeserializationSchema<>(Customer.class))
                .build();

        DataStreamSource<Customer> dataStreamSource =
                env.fromSource(source, WatermarkStrategy.noWatermarks(), "Kafka Source");

        dataStreamSource.print();

        JdbcSink<Customer> customerJdbcSink = new JdbcSink<>(
                DeliveryGuarantee.AT_LEAST_ONCE,
                new SimpleJdbcConnectionProvider(
                        new JdbcConnectionOptions.JdbcConnectionOptionsBuilder()
                        .withUrl("jdbc:mysql://106.54.174.109:13306/my_cdc")
                        .withDriverName("com.mysql.jdbc.Driver")
                        .withUsername("root")
                        .withPassword("mysql@123456")
                        .build()
                ),
                JdbcExecutionOptions.builder()
                        .withMaxRetries(1)
                        .withBatchSize(10)
                        .withBatchIntervalMs(5000)
                        .build(),
                JdbcExactlyOnceOptions.defaults(),
                new JdbcQueryStatement<>() {
                    @Override
                    public String query() {
                        return "insert into customers_kafka (featuresId, siteSn, captureTime) values (?, ?, ?)";
                    }

                    @Override
                    public void statement(PreparedStatement ps, Customer customer) throws SQLException {
                        ps.setString(1, customer.getFeaturesId());
                        ps.setString(2, customer.getSiteSn());
                        ps.setString(3, Long.toString(customer.getCaptureTime()));
                    }
                }
        );

        dataStreamSource.sinkTo(customerJdbcSink);

        env.execute();
    }

    private static void jdbcSource() throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);

        JdbcSource<Customer> jdbcSource = JdbcSource.<Customer>builder()
                .setDBUrl("jdbc:mysql://106.54.174.109:13306/my_cdc")
                .setDriverName("com.mysql.jdbc.Driver")
                .setUsername("root")
                .setPassword("mysql@123456")
                .setSql("select featuresId, siteSn, captureTime from spend_report limit 100")
                .setResultExtractor((ResultExtractor<Customer>) resultSet -> {
                    String featuresId = resultSet.getString("featuresId");
                    String siteSn = resultSet.getString("siteSn");
                    long captureTime = resultSet.getLong("captureTime");

                    Customer customer = new Customer();
                    customer.setFeaturesId(featuresId);
                    customer.setSiteSn(siteSn);
                    customer.setCaptureTime(captureTime);
                    return customer;
                })
                .setTypeInformation(TypeInformation.of(Customer.class))
                .build();

        DataStreamSource<Customer> dataStreamSource =
                env.fromSource(jdbcSource, WatermarkStrategy.noWatermarks(), "jdbc source");

        dataStreamSource.print();

        env.execute();
    }

}
